#!/usr/bin/env python3
"""
EchoKit ROS2 Bridge - EchoKit服务器到ROS2节点的通信桥接
实现EchoKit Server与ROS2视觉交互节点之间的双向通信
"""

import asyncio
import json
import threading
import time
from typing import Dict, Any, Optional, Callable
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, DurabilityPolicy
from std_msgs.msg import String
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import numpy as np

from .enhanced_websocket_handler import EnhancedWebSocketHandler
from .vision_config_manager import VisionConfigManager


class EchoKitROS2Bridge(Node):
    """EchoKit服务器到ROS2的通信桥接"""
    
    def __init__(self, websocket_url: str, config_manager: Optional[VisionConfigManager] = None):
        super().__init__('echokit_ros2_bridge')
        
        # 配置管理器
        self.config_manager = config_manager or VisionConfigManager()
        
        # WebSocket处理器
        self.websocket_handler = EnhancedWebSocketHandler(
            url=websocket_url,
            logger=self.get_logger()
        )
        
        # CV Bridge for image conversion
        self.cv_bridge = CvBridge()
        
        # QoS配置
        self.qos_profile = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            durability=DurabilityPolicy.TRANSIENT_LOCAL,
            depth=10
        )
        
        # ROS2发布者
        self.vision_query_publisher = self.create_publisher(
            String, '/vision/query', self.qos_profile
        )
        self.vision_response_publisher = self.create_publisher(
            String, '/vision/response', self.qos_profile
        )
        self.camera_status_publisher = self.create_publisher(
            String, '/camera/status', self.qos_profile
        )
        
        # ROS2订阅者
        self.vision_analysis_subscriber = self.create_subscription(
            String, '/vision_analysis', self.vision_analysis_callback, self.qos_profile
        )
        self.current_frame_subscriber = self.create_subscription(
            Image, '/vision/current_frame', self.current_frame_callback, self.qos_profile
        )
        self.voice_text_subscriber = self.create_subscription(
            String, '/voice_text', self.voice_text_callback, self.qos_profile
        )
        
        # 状态变量
        self.current_frame = None
        self.frame_timestamp = None
        self.connected_to_echokit = False
        self.pending_queries = {}  # query_id -> query_data
        
        # 设置WebSocket回调
        self._setup_websocket_callbacks()
        
        # 启动WebSocket连接
        self.websocket_handler.start()
        
        self.get_logger().info('EchoKit ROS2 Bridge initialized')
        self.get_logger().info(f'WebSocket URL: {websocket_url}')
    
    def _setup_websocket_callbacks(self):
        """设置WebSocket回调函数"""
        self.websocket_handler.set_connection_callback(self._on_websocket_connection_changed)
        self.websocket_handler.set_vision_query_callback(self._on_vision_query_received)
        self.websocket_handler.set_vision_response_callback(self._on_vision_response_received)
        self.websocket_handler.set_camera_status_callback(self._on_camera_status_received)
        self.websocket_handler.set_error_callback(self._on_websocket_error)
    
    def _on_websocket_connection_changed(self, connected: bool):
        """WebSocket连接状态变化回调"""
        self.connected_to_echokit = connected
        status = "connected" if connected else "disconnected"
        self.get_logger().info(f'EchoKit Server connection: {status}')
        
        # 发布连接状态
        status_msg = String()
        status_msg.data = json.dumps({
            'type': 'echokit_connection',
            'connected': connected,
            'timestamp': int(time.time() * 1000)
        })
        self.camera_status_publisher.publish(status_msg)
    
    def _on_vision_query_received(self, query_data: Dict[str, Any]):
        """接收到视觉查询回调"""
        self.get_logger().info(f'Received vision query from EchoKit: {query_data.get("query_text", "")}')
        
        # 存储查询信息
        query_id = query_data.get('query_id', '')
        self.pending_queries[query_id] = query_data
        
        # 发布到ROS2话题
        query_msg = String()
        query_msg.data = json.dumps(query_data)
        self.vision_query_publisher.publish(query_msg)
        
        # 如果有当前帧，触发视觉分析
        if self.current_frame is not None:
            self._trigger_vision_analysis(query_data)
    
    def _on_vision_response_received(self, response_data: Dict[str, Any]):
        """接收到视觉响应回调"""
        self.get_logger().info(f'Received vision response from EchoKit: {response_data.get("response_text", "")}')
        
        # 发布到ROS2话题
        response_msg = String()
        response_msg.data = json.dumps(response_data)
        self.vision_response_publisher.publish(response_msg)
    
    def _on_camera_status_received(self, status_data: Dict[str, Any]):
        """接收到摄像头状态回调"""
        self.get_logger().info(f'Received camera status from EchoKit: connected={status_data.get("connected", False)}')
        
        # 发布到ROS2话题
        status_msg = String()
        status_msg.data = json.dumps(status_data)
        self.camera_status_publisher.publish(status_msg)
    
    def _on_websocket_error(self, error: str):
        """WebSocket错误回调"""
        self.get_logger().error(f'EchoKit WebSocket error: {error}')
        
        # 发布错误状态
        error_msg = String()
        error_msg.data = json.dumps({
            'type': 'echokit_error',
            'error': error,
            'timestamp': int(time.time() * 1000)
        })
        self.camera_status_publisher.publish(error_msg)
    
    def vision_analysis_callback(self, msg: String):
        """处理视觉分析结果"""
        try:
            analysis_data = json.loads(msg.data)
            self.get_logger().info(f'Received vision analysis: {analysis_data}')
            
            # 检查是否有对应的查询
            query_id = analysis_data.get('query_id', '')
            if query_id in self.pending_queries:
                # 发送响应到EchoKit Server
                response_text = analysis_data.get('analysis_result', '')
                confidence = analysis_data.get('confidence_score', 1.0)
                processing_time = analysis_data.get('processing_time_ms', 0)
                model_used = analysis_data.get('model_used', 'unknown')
                
                success = self.websocket_handler.send_vision_response(
                    query_id, response_text, confidence, processing_time, model_used
                )
                
                if success:
                    self.get_logger().info(f'Sent vision response to EchoKit for query {query_id}')
                    # 清理已处理的查询
                    del self.pending_queries[query_id]
                else:
                    self.get_logger().error(f'Failed to send vision response to EchoKit for query {query_id}')
            else:
                self.get_logger().warn(f'Received analysis for unknown query: {query_id}')
                
        except json.JSONDecodeError as e:
            self.get_logger().error(f'Failed to parse vision analysis data: {e}')
        except Exception as e:
            self.get_logger().error(f'Error processing vision analysis: {e}')
    
    def current_frame_callback(self, msg: Image):
        """处理当前帧更新"""
        try:
            # 转换图像格式
            self.current_frame = self.cv_bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
            self.frame_timestamp = time.time()
            
            self.get_logger().debug(f'Updated current frame: {self.current_frame.shape}')
            
        except Exception as e:
            self.get_logger().error(f'Error processing current frame: {e}')
    
    def voice_text_callback(self, msg: String):
        """处理语音文本输入"""
        try:
            voice_text = msg.data
            self.get_logger().info(f'Received voice text: {voice_text}')
            
            # 检查是否是视觉查询
            if self.websocket_handler.message_handler.is_vision_query(voice_text):
                self.get_logger().info('Detected vision query in voice text')
                
                # 创建视觉查询
                query_data = {
                    'query_id': str(int(time.time() * 1000)),
                    'query_text': voice_text,
                    'timestamp': int(time.time() * 1000),
                    'source': 'voice'
                }
                
                # 如果有当前帧，附加图像数据
                if self.current_frame is not None:
                    # 将图像编码为JPEG字节
                    import cv2
                    _, encoded_image = cv2.imencode('.jpg', self.current_frame)
                    query_data['image_data'] = encoded_image.tobytes()
                
                # 发送到EchoKit Server
                success = self.websocket_handler.send_vision_query(
                    query_data['query_text'], 
                    query_data.get('image_data')
                )
                
                if success:
                    self.get_logger().info('Sent vision query to EchoKit Server')
                    self.pending_queries[query_data['query_id']] = query_data
                else:
                    self.get_logger().error('Failed to send vision query to EchoKit Server')
            
        except Exception as e:
            self.get_logger().error(f'Error processing voice text: {e}')
    
    def _trigger_vision_analysis(self, query_data: Dict[str, Any]):
        """触发视觉分析"""
        if self.current_frame is None:
            self.get_logger().warn('No current frame available for vision analysis')
            return
        
        try:
            # 创建视觉查询消息
            analysis_request = {
                'query_id': query_data.get('query_id', ''),
                'query_text': query_data.get('query_text', ''),
                'timestamp': int(time.time() * 1000),
                'frame_timestamp': self.frame_timestamp
            }
            
            # 发布视觉查询请求
            query_msg = String()
            query_msg.data = json.dumps(analysis_request)
            self.vision_query_publisher.publish(query_msg)
            
            self.get_logger().info(f'Triggered vision analysis for query: {query_data.get("query_id", "")}')
            
        except Exception as e:
            self.get_logger().error(f'Error triggering vision analysis: {e}')
    
    def send_camera_status_to_echokit(self, connected: bool, resolution: str = '', 
                                    fps: int = 0, format_type: str = '',
                                    stream_quality: float = 0.0):
        """发送摄像头状态到EchoKit Server"""
        if not self.connected_to_echokit:
            self.get_logger().warn('Not connected to EchoKit Server, cannot send camera status')
            return False
        
        success = self.websocket_handler.send_camera_status(
            connected, resolution, fps, format_type, stream_quality
        )
        
        if success:
            self.get_logger().info('Sent camera status to EchoKit Server')
        else:
            self.get_logger().error('Failed to send camera status to EchoKit Server')
        
        return success
    
    def get_bridge_status(self) -> Dict[str, Any]:
        """获取桥接状态"""
        return {
            'connected_to_echokit': self.connected_to_echokit,
            'has_current_frame': self.current_frame is not None,
            'frame_timestamp': self.frame_timestamp,
            'pending_queries': len(self.pending_queries),
            'websocket_stats': self.websocket_handler.get_stats()
        }
    
    def shutdown(self):
        """关闭桥接"""
        self.get_logger().info('Shutting down EchoKit ROS2 Bridge')
        
        # 停止WebSocket连接
        self.websocket_handler.stop()
        
        # 清理资源
        self.pending_queries.clear()
        self.current_frame = None


def main(args=None):
    """主函数"""
    rclpy.init(args=args)
    
    # 从参数获取WebSocket URL
    import sys
    websocket_url = 'ws://<ECHO_SERVER>:8060/ws/robot'
    
    if len(sys.argv) > 1:
        websocket_url = sys.argv[1]
    
    try:
        # 创建配置管理器
        config_manager = VisionConfigManager()
        
        # 创建桥接节点
        bridge = EchoKitROS2Bridge(websocket_url, config_manager)
        
        # 运行节点
        rclpy.spin(bridge)
        
    except KeyboardInterrupt:
        print('EchoKit ROS2 Bridge interrupted by user')
    except Exception as e:
        print(f'Error in EchoKit ROS2 Bridge: {e}')
    finally:
        if 'bridge' in locals():
            bridge.shutdown()
        rclpy.shutdown()


if __name__ == '__main__':
    main()